import json
import math
import yaml
import os
import shutil


def replace_nan_by_none(dict):
    dict = {k: (None if math.isnan(v) else v) if isinstance(v, float) else v for k, v in dict.items()}
    return dict


def dict_list_to_json(dict_list, fname):
    '''
    Dump a list of Python dictionaries to JSON file

    Args:
        dict_list (list of dict): List of Python dictionaries
        fname (str): Output file name

    '''
    with open(fname, 'a') as f:
        for dict in dict_list:
            json.dump(dict, f)
            f.write('\n')


def dict_to_json(dict, fname):
    '''
    Dump a Python dictionary to JSON file
    Args:
        dict (dict): Python dictionary
        fname (str): Output file name
    '''
    with open(fname, 'a') as f:
        json.dump(dict, f)
        f.write('\n')


def json_to_dict_list(fname):
    dict_list = []
    epoch_set = set()
    with open(fname) as f:
        lines = f.readlines()
        for line in lines:
            line = line.rstrip()
            dict = json.loads(line)
            if dict.get('epoch') not in epoch_set:
                dict_list.append(dict)
            epoch_set.add(dict.get('epoch'))
    return dict_list


def dict_list_to_tb(dict_list, writer):
    for dict in dict_list:
        assert 'epoch' in dict, 'Key epoch must exist in stats dict'
        dict_to_tb(dict, writer, dict['epoch'])


def makedirs(dir):
    os.makedirs(dir, exist_ok=True)


def makedirs_rm_exist(dir):
    '''
    Make a directory, remove any existing data.
    Args:
        dir (str): The directory to be created.
    Returns:
    '''
    if os.path.isdir(dir):
        shutil.rmtree(dir)
    os.makedirs(dir, exist_ok=True)
    

def dict_to_tb(dict, writer, epoch):
    '''
    Add a dictionary of statistics to a Tensorboard writer
    Args:
        dict (dict): Statistics of experiments, the keys are attribute names,
        the values are the attribute values
        writer: Tensorboard writer object
        epoch (int): The current epoch
    '''
    for key in dict:
        writer.add_scalar(key, dict[key], epoch)


def dict_to_neptune(name, dict, writer, epoch):
    for key in dict:
        writer[f'{name}/{key}'].append(dict[key], step=epoch)


def load_config(file_str: str = 'config.yaml') -> None:
    with open(file_str) as f:
        yaml.safe_load(f)
